"""
Prompt优化器 - 负责合并和去重多层Prompt
"""
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from enum import Enum
import re
import hashlib


class PromptLayer(Enum):
    """Prompt层级枚举"""
    TEMPLATE = "template"      # 模板层
    SYSTEM = "system"         # 系统层  
    SMART = "smart"           # 智能层
    FRONTEND = "frontend"     # 界面层
    USER = "user"            # 用户层


@dataclass
class PromptSegment:
    """Prompt片段"""
    layer: PromptLayer
    type: str  # role, task, context, constraint, requirement
    content: str
    priority: int = 0
    metadata: Dict[str, Any] = None


class PromptOptimizer:
    """Prompt优化器"""
    
    def __init__(self):
        self.deduplication_threshold = 0.8  # 去重相似度阈值
        self.max_total_length = 4000  # 最大总长度
        
    def optimize_prompts(self, prompt_layers: Dict[PromptLayer, str]) -> str:
        """
        优化多层Prompt，去重合并
        
        Args:
            prompt_layers: 各层Prompt内容
            
        Returns:
            优化后的合并Prompt
        """
        # 1. 解析各层Prompt为结构化片段
        segments = self._parse_prompt_layers(prompt_layers)
        
        # 2. 去重相似内容
        deduplicated_segments = self._deduplicate_segments(segments)
        
        # 3. 按类型重新组织
        organized_segments = self._organize_by_type(deduplicated_segments)
        
        # 4. 合并生成最终Prompt
        final_prompt = self._merge_segments(organized_segments)
        
        # 5. 长度控制
        optimized_prompt = self._control_length(final_prompt)
        
        return optimized_prompt
    
    def _parse_prompt_layers(self, prompt_layers: Dict[PromptLayer, str]) -> List[PromptSegment]:
        """解析各层Prompt为结构化片段"""
        segments = []
        
        for layer, content in prompt_layers.items():
            if not content:
                continue
                
            # 解析不同类型的内容
            parsed_segments = self._parse_content_by_layer(layer, content)
            segments.extend(parsed_segments)
            
        return segments
    
    def _parse_content_by_layer(self, layer: PromptLayer, content: str) -> List[PromptSegment]:
        """根据层级解析内容"""
        segments = []
        
        if layer == PromptLayer.TEMPLATE:
            # 模板层主要是角色定义
            segments.append(PromptSegment(
                layer=layer,
                type="role",
                content=content,
                priority=5
            ))
            
        elif layer == PromptLayer.SYSTEM:
            # 系统层包含角色、上下文、约束
            role_content = self._extract_role_definition(content)
            context_content = self._extract_context_info(content)
            constraint_content = self._extract_constraints(content)
            
            if role_content:
                segments.append(PromptSegment(layer=layer, type="role", content=role_content, priority=4))
            if context_content:
                segments.append(PromptSegment(layer=layer, type="context", content=context_content, priority=3))
            if constraint_content:
                segments.append(PromptSegment(layer=layer, type="constraint", content=constraint_content, priority=3))
                
        elif layer == PromptLayer.SMART:
            # 智能层主要是任务定义
            task_content = self._extract_task_definition(content)
            requirement_content = self._extract_requirements(content)
            
            if task_content:
                segments.append(PromptSegment(layer=layer, type="task", content=task_content, priority=4))
            if requirement_content:
                segments.append(PromptSegment(layer=layer, type="requirement", content=requirement_content, priority=2))
                
        elif layer == PromptLayer.FRONTEND:
            # 界面层主要是要求
            requirement_content = self._extract_requirements(content)
            if requirement_content:
                segments.append(PromptSegment(layer=layer, type="requirement", content=requirement_content, priority=1))
                
        elif layer == PromptLayer.USER:
            # 用户层是个性化要求
            segments.append(PromptSegment(
                layer=layer,
                type="user_requirement",
                content=content,
                priority=5
            ))
            
        return segments
    
    def _extract_role_definition(self, content: str) -> str:
        """提取角色定义"""
        patterns = [
            r"你是.*?(?=\n|$)",
            r"作为.*?(?=\n|$)",
            r"扮演.*?(?=\n|$)"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, content)
            if match:
                return match.group(0)
        return ""
    
    def _extract_context_info(self, content: str) -> str:
        """提取上下文信息"""
        context_sections = []
        
        # 提取热点信息
        hot_info_match = re.search(r"当前热点信息：.*?(?=\n\n|\n[^-]|$)", content, re.DOTALL)
        if hot_info_match:
            context_sections.append(hot_info_match.group(0))
            
        # 提取评论分析
        comment_match = re.search(r"评论分析：.*?(?=\n\n|\n[^-]|$)", content, re.DOTALL)
        if comment_match:
            context_sections.append(comment_match.group(0))
            
        return "\n\n".join(context_sections)
    
    def _extract_constraints(self, content: str) -> str:
        """提取约束条件"""
        constraint_patterns = [
            r"避免.*?(?=\n|$)",
            r"禁止.*?(?=\n|$)",
            r"不能.*?(?=\n|$)",
            r"遵循.*?(?=\n|$)"
        ]
        
        constraints = []
        for pattern in constraint_patterns:
            matches = re.findall(pattern, content)
            constraints.extend(matches)
            
        return "\n".join(constraints) if constraints else ""
    
    def _extract_task_definition(self, content: str) -> str:
        """提取任务定义"""
        task_patterns = [
            r"针对.*?撰写.*?(?=\n|$)",
            r"请.*?(?=\n|$)",
            r"任务：.*?(?=\n|$)"
        ]
        
        for pattern in task_patterns:
            match = re.search(pattern, content)
            if match:
                return match.group(0)
        return ""
    
    def _extract_requirements(self, content: str) -> str:
        """提取创作要求"""
        # 提取编号列表要求
        numbered_requirements = re.findall(r"\d+\.\s*[^.\n]+", content)
        
        # 提取其他要求格式
        bullet_requirements = re.findall(r"[-•]\s*[^.\n]+", content)
        
        all_requirements = numbered_requirements + bullet_requirements
        
        # 去重相似要求
        unique_requirements = self._deduplicate_requirements(all_requirements)
        
        return "\n".join([f"{i+1}. {req.strip()}" for i, req in enumerate(unique_requirements)])
    
    def _deduplicate_requirements(self, requirements: List[str]) -> List[str]:
        """去重相似的要求"""
        unique_requirements = []
        
        for req in requirements:
            is_duplicate = False
            req_clean = self._normalize_text(req)
            
            for existing in unique_requirements:
                existing_clean = self._normalize_text(existing)
                similarity = self._calculate_similarity(req_clean, existing_clean)
                
                if similarity > self.deduplication_threshold:
                    is_duplicate = True
                    break
                    
            if not is_duplicate:
                unique_requirements.append(req)
                
        return unique_requirements
    
    def _deduplicate_segments(self, segments: List[PromptSegment]) -> List[PromptSegment]:
        """去重相似的片段"""
        deduplicated = []
        
        # 按类型分组处理
        segments_by_type = {}
        for segment in segments:
            if segment.type not in segments_by_type:
                segments_by_type[segment.type] = []
            segments_by_type[segment.type].append(segment)
        
        # 每个类型内部去重
        for segment_type, type_segments in segments_by_type.items():
            # 按优先级排序
            type_segments.sort(key=lambda x: x.priority, reverse=True)
            
            unique_segments = []
            for segment in type_segments:
                is_duplicate = False
                segment_clean = self._normalize_text(segment.content)
                
                for existing in unique_segments:
                    existing_clean = self._normalize_text(existing.content)
                    similarity = self._calculate_similarity(segment_clean, existing_clean)
                    
                    if similarity > self.deduplication_threshold:
                        # 如果相似，保留优先级更高的
                        if segment.priority > existing.priority:
                            unique_segments.remove(existing)
                            unique_segments.append(segment)
                        is_duplicate = True
                        break
                        
                if not is_duplicate:
                    unique_segments.append(segment)
                    
            deduplicated.extend(unique_segments)
            
        return deduplicated
    
    def _organize_by_type(self, segments: List[PromptSegment]) -> Dict[str, List[PromptSegment]]:
        """按类型重新组织片段"""
        organized = {
            "role": [],
            "task": [],
            "context": [],
            "constraint": [],
            "requirement": [],
            "user_requirement": []
        }
        
        for segment in segments:
            if segment.type in organized:
                organized[segment.type].append(segment)
                
        return organized
    
    def _merge_segments(self, organized_segments: Dict[str, List[PromptSegment]]) -> str:
        """合并片段生成最终Prompt"""
        prompt_parts = []
        
        # 1. 角色定义
        if organized_segments["role"]:
            role_content = self._merge_role_segments(organized_segments["role"])
            prompt_parts.append(role_content)
        
        # 2. 任务定义
        if organized_segments["task"]:
            task_content = self._merge_task_segments(organized_segments["task"])
            prompt_parts.append(f"任务：{task_content}")
        
        # 3. 上下文信息
        if organized_segments["context"]:
            context_content = self._merge_context_segments(organized_segments["context"])
            prompt_parts.append(f"上下文：\n{context_content}")
        
        # 4. 创作要求
        requirements = organized_segments["requirement"] + organized_segments["user_requirement"]
        if requirements:
            req_content = self._merge_requirement_segments(requirements)
            prompt_parts.append(f"创作要求：\n{req_content}")
        
        # 5. 约束条件
        if organized_segments["constraint"]:
            constraint_content = self._merge_constraint_segments(organized_segments["constraint"])
            prompt_parts.append(f"注意事项：\n{constraint_content}")
        
        return "\n\n".join(prompt_parts)
    
    def _merge_role_segments(self, role_segments: List[PromptSegment]) -> str:
        """合并角色片段"""
        # 选择优先级最高的角色定义
        role_segments.sort(key=lambda x: x.priority, reverse=True)
        return role_segments[0].content if role_segments else ""
    
    def _merge_task_segments(self, task_segments: List[PromptSegment]) -> str:
        """合并任务片段"""
        # 选择最具体的任务定义
        task_segments.sort(key=lambda x: len(x.content), reverse=True)
        return task_segments[0].content if task_segments else ""
    
    def _merge_context_segments(self, context_segments: List[PromptSegment]) -> str:
        """合并上下文片段"""
        context_parts = []
        for segment in context_segments:
            if segment.content.strip():
                context_parts.append(segment.content.strip())
        return "\n".join(context_parts)
    
    def _merge_requirement_segments(self, requirement_segments: List[PromptSegment]) -> str:
        """合并要求片段"""
        all_requirements = []
        
        for segment in requirement_segments:
            # 解析要求列表
            requirements = re.findall(r"\d+\.\s*([^.\n]+)", segment.content)
            if not requirements:
                requirements = [segment.content.strip()]
            all_requirements.extend(requirements)
        
        # 去重并重新编号
        unique_requirements = self._deduplicate_requirements(all_requirements)
        return "\n".join([f"{i+1}. {req.strip()}" for i, req in enumerate(unique_requirements)])
    
    def _merge_constraint_segments(self, constraint_segments: List[PromptSegment]) -> str:
        """合并约束片段"""
        constraints = []
        for segment in constraint_segments:
            if segment.content.strip():
                constraints.append(segment.content.strip())
        return "\n".join(constraints)
    
    def _control_length(self, prompt: str) -> str:
        """控制Prompt长度"""
        if len(prompt) <= self.max_total_length:
            return prompt
        
        # 如果超长，按优先级截断
        lines = prompt.split('\n')
        truncated_lines = []
        current_length = 0
        
        for line in lines:
            if current_length + len(line) <= self.max_total_length:
                truncated_lines.append(line)
                current_length += len(line)
            else:
                break
        
        return '\n'.join(truncated_lines)
    
    def _normalize_text(self, text: str) -> str:
        """标准化文本用于比较"""
        # 移除标点符号和空格，转小写
        normalized = re.sub(r'[^\w\u4e00-\u9fff]', '', text.lower())
        return normalized
    
    def _calculate_similarity(self, text1: str, text2: str) -> float:
        """计算文本相似度"""
        if not text1 or not text2:
            return 0.0
        
        # 使用简单的字符级相似度
        set1 = set(text1)
        set2 = set(text2)
        
        intersection = len(set1.intersection(set2))
        union = len(set1.union(set2))
        
        return intersection / union if union > 0 else 0.0


# 使用示例
def optimize_hot_topic_prompts(
    template_prompt: str = "",
    system_prompt: str = "",
    smart_prompt: str = "",
    frontend_prompt: str = "",
    user_input: str = ""
) -> str:
    """优化热点话题的多层Prompt"""
    
    optimizer = PromptOptimizer()
    
    prompt_layers = {
        PromptLayer.TEMPLATE: template_prompt,
        PromptLayer.SYSTEM: system_prompt,
        PromptLayer.SMART: smart_prompt,
        PromptLayer.FRONTEND: frontend_prompt,
        PromptLayer.USER: user_input
    }
    
    return optimizer.optimize_prompts(prompt_layers)
